package com.hzsparrow.ai.config;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.ai.deepseek.api.DeepSeekApi;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.deepseek.DeepSeekChatOptions;
import org.springframework.stereotype.Component;

import com.hzsparrow.ai.entity.ModelConfig;

import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 聊天模型管理器
 * 用于动态创建和管理不同类型的聊天模型
 */
@Slf4j
@Component
public class ChatModelManager {
    
    private final Map<Long, ChatModel> chatModelCache = new ConcurrentHashMap<>();
    
    /**
     * 根据模型配置创建聊天模型
     * 
     * @param modelConfig 模型配置
     * @return 聊天模型实例
     */
    public ChatModel createChatModel(ModelConfig modelConfig) {
        if (modelConfig == null) {
            throw new IllegalArgumentException("模型配置不能为空");
        }
        
        if (!"chat".equals(modelConfig.getModelType())) {
            throw new IllegalArgumentException("模型类型必须是chat");
        }
        
        try {
            ChatModel chatModel = switch (modelConfig.getCategory().toLowerCase()) {
                case "ollama" -> createOllamaChatModel(modelConfig);
                case "openai" -> createOpenAiChatModel(modelConfig);
                case "deepseek" -> createDeepSeekChatModel(modelConfig);
                default -> throw new IllegalArgumentException("不支持的模型分类: " + modelConfig.getCategory());
            };
            
            log.info("创建聊天模型成功: {} - {}", modelConfig.getAlias(), modelConfig.getCategory());
            return chatModel;
        } catch (Exception e) {
            log.error("创建聊天模型失败: {} - {}", modelConfig.getAlias(), e.getMessage());
            throw new RuntimeException("创建聊天模型失败", e);
        }
    }
    
    /**
     * 获取或创建聊天模型（带缓存）
     * 
     * @param modelConfig 模型配置
     * @return 聊天模型实例
     */
    public ChatModel getOrCreateChatModel(ModelConfig modelConfig) {
        return chatModelCache.computeIfAbsent(modelConfig.getId(), 
            k -> createChatModel(modelConfig));
    }
    
    /**
     * 清除指定模型的缓存
     * 
     * @param modelId 模型ID
     */
    public void clearCache(Long modelId) {
        chatModelCache.remove(modelId);
        log.info("清除聊天模型缓存: {}", modelId);
    }
    
    /**
     * 清除所有缓存
     */
    public void clearAllCache() {
        chatModelCache.clear();
        log.info("清除所有聊天模型缓存");
    }
    
    /**
     * 创建Ollama聊天模型
     */
    private ChatModel createOllamaChatModel(ModelConfig modelConfig) {
        OllamaApi ollamaApi = OllamaApi.builder()
            .baseUrl(modelConfig.getBasePath())
            .build();
        
        return OllamaChatModel.builder()
            .ollamaApi(ollamaApi)
            .defaultOptions(OllamaOptions.builder()
                .model(modelConfig.getModelName())
                .build())
            .build();
    }
    
    /**
     * 创建OpenAI聊天模型
     */
    private ChatModel createOpenAiChatModel(ModelConfig modelConfig) {
        OpenAiApi openAiApi = OpenAiApi.builder()
            .baseUrl(modelConfig.getBasePath())
            .apiKey(modelConfig.getApiKey())
            .build();
        
        return OpenAiChatModel.builder()
            .openAiApi(openAiApi)
            .defaultOptions(OpenAiChatOptions.builder()
                .model(modelConfig.getModelName())
                .build())
            .build();
    }
    
    /**
     * 创建DeepSeek聊天模型
     */
    private ChatModel createDeepSeekChatModel(ModelConfig modelConfig) {
        DeepSeekApi deepSeekApi = DeepSeekApi.builder()
            .baseUrl(modelConfig.getBasePath())
            .apiKey(modelConfig.getApiKey())
            .build();
        
        return DeepSeekChatModel.builder()
            .deepSeekApi(deepSeekApi)
            .defaultOptions(DeepSeekChatOptions.builder()
                .model(modelConfig.getModelName())
                .build())
            .build();
    }
} 